import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import datasets  # 导入经典数据集加载模块

# 加载mnist数据集
(x, y), (x_test, y_test) = datasets.mnist.load_data()
print(x.shape)
print(x_test.shape)

# 将数据转换为Dataset对象，才能利用TensorFlow提供的各种便捷功能
train_db = tf.data.Dataset.from_tensor_slices((x, y))
# 随机打撒样本，并不会打乱样本和标间的顺序
# 10000为buffer_size(缓冲池的大小，一般设置为一个较大的常数即可)
train_db = train_db.shuffle(10000)

# 批训练设置Batch Size
train_db = train_db.batch(128)


# 预处理
def preprocess(x, y):
    # 调用此函数是会自动传入x,y对象，shape为[b,28,28],[b]
    x = tf.constant(x, dtype=tf.float32) / 255.
    x = tf.reshape(x, [-1, 28 * 28])  # 打平
    y = tf.constant(y, dtype=tf.float32) / 255.
    y = tf.one_hot(y, depth=10)  # 独热编码
    return x, y

# 循环训练
# for step,(x,y) in enumerate(train_db):

